package com.oliver.proxy.midjourney.application.service;

import com.alibaba.fastjson2.JSONObject;
import com.oliver.proxy.midjourney.domain.entity.Prod;
import com.oliver.proxy.midjourney.domain.entity.TaskList;
import com.oliver.proxy.midjourney.infrastructure.mapper.ProdMapper;
import com.oliver.proxy.midjourney.infrastructure.mapper.TaskListMapper;
import io.github.yuanbaobaoo.dify.client.DifyClientBuilder;
import io.github.yuanbaobaoo.dify.client.IDifyWorkFlowClient;
import io.github.yuanbaobaoo.dify.client.params.ParamMessage;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
@Service
public class DifyProcessingService {
    private final ProdMapper prodMapper;
    private final TaskListMapper taskListMapper;
    private final ExecutorService executor;
    private final AtomicBoolean isProcessing = new AtomicBoolean(false);
    private static final int THREAD_COUNT = 4; // 使用4个线程并行处理
    private final List<IDifyWorkFlowClient> difyClients = new ArrayList<>();

    @Autowired
    public DifyProcessingService(ProdMapper prodMapper, TaskListMapper taskListMapper) {
        this.prodMapper = prodMapper;
        this.taskListMapper = taskListMapper;
        
        // 创建多个Dify客户端实例，每个线程一个
        for (int i = 0; i < THREAD_COUNT; i++) {
            difyClients.add(DifyClientBuilder.create()
                    .apiKey("app-XAJNiJ3RArnqXYgUj07VUpyi")
                    .baseUrl("http://localhost/v1")
                    .buildWorkFlow());
        }
        
        // 创建固定大小的线程池
        this.executor = Executors.newFixedThreadPool(THREAD_COUNT);
    }

    @PostConstruct
    public void init() {
        // 在后台线程中启动处理
        CompletableFuture.runAsync(this::processDifyRequestsInBackground);
    }

    // 定时检查是否需要重新启动处理
    @Scheduled(fixedRate = 6000) // 每分钟检查一次
    public void checkAndRestartProcessing() {
        if (!isProcessing.get()) {
            log.info("Restarting Dify processing...");
            CompletableFuture.runAsync(this::processDifyRequestsInBackground);
        }
    }

    private void processDifyRequestsInBackground() {
        try {
            if (!isProcessing.compareAndSet(false, true)) {
                return; // 已经有线程在处理
            }
            
            log.info("Starting Dify processing in background thread");
            
            while (true) {
                try {
                    // 获取所有需要处理的记录（title为空的记录）
                    List<Prod> unprocessedProds = prodMapper.selectProdsWithTitle();
                    if (unprocessedProds.isEmpty()) {
                        Thread.sleep(2000);
                        continue;
                    }

                    log.info("Found {} prods to process", unprocessedProds.size());
                    
                    // 分批处理记录
                    processProdsInParallel(unprocessedProds);
                } catch (Exception e) {
                    log.error("Error processing batch of prods: ", e);
                    // 添加短暂休眠，防止因异常导致的CPU高负载循环
                    Thread.sleep(5000);
                }
            }
        } catch (Exception e) {
            log.error("Error in Dify background processing", e);
        } finally {
            isProcessing.set(false);
        }
    }

    private void processProdsInParallel(List<Prod> prods) throws InterruptedException {
        // 将记录分成多个批次
        List<List<Prod>> batches = new ArrayList<>();
        for (int i = 0; i < THREAD_COUNT; i++) {
            batches.add(new ArrayList<>());
        }
        
        // 将记录均匀分配到各个批次
        for (int i = 0; i < prods.size(); i++) {
            batches.get(i % THREAD_COUNT).add(prods.get(i));
        }
        
        // 创建任务列表
        List<CompletableFuture<Void>> futures = new ArrayList<>();
        
        // 为每个批次创建一个任务
        for (int i = 0; i < THREAD_COUNT; i++) {
            final int threadIndex = i;
            List<Prod> batch = batches.get(i);
            
            if (!batch.isEmpty()) {
                CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
                    processBatch(batch, threadIndex);
                }, executor);
                
                futures.add(future);
            }
        }
        
        // 等待所有任务完成
        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
    }
    
    private void processBatch(List<Prod> batch, int threadIndex) {
        IDifyWorkFlowClient client = difyClients.get(threadIndex);
        log.info("Thread {} processing {} records", threadIndex, batch.size());
        
        for (Prod prod : batch) {
            try {
                processSingleProd(prod, client);
                Thread.sleep(1000); // 每次处理间隔1秒
            } catch (Exception e) {
                log.error("Thread {} error processing prod: {}", threadIndex, prod.getDataId(), e);
                // 不再抛出异常，继续处理下一条记录
            }
        }
    }

    @Transactional
    public void processSingleProd(Prod prod, IDifyWorkFlowClient client) {
        try {
            // 获取对应的TaskList记录
            TaskList taskList = taskListMapper.selectById(prod.getTaskId());
            if (taskList == null) {
                log.warn("TaskList not found for prod: {}", prod.getDataId());
                return;
            }

            // 创建消息
            ParamMessage message = ParamMessage
                    .builder()
                    .user("x@oliverding.cn")
                    .inputs(new HashMap<>() {{
                        put("prompt", prod.getPrompt());
                        put("prod", taskList.getTaskCategory());
                    }})
                    .build();

            // 调用Dify API - 添加超时处理
            JSONObject result;
            try {
                result = client.runBlocking(message);
            } catch (Exception e) {
                log.error("Dify API call timed out or failed for prod: {}", prod.getDataId(), e);
                return; // 返回而不抛出异常
            }

            // 更新prod记录
            if (null != result) {
                JSONObject data = result.getJSONObject("data");
                if (data != null && "succeeded".equals(data.getString("status"))) {
                    String title = data.getJSONObject("outputs").getString("text");
                    title = truncateTitle(title, 240);
                    prod.setTitle(title);
                    prodMapper.updateTitle(prod);
                    log.info("Successfully updated title for prod: {}", prod.getDataId());
                }
            }
        } catch (Exception e) {
            log.error("Error processing prod with Dify: {}", prod.getDataId(), e);
            // 移除抛出异常，避免整个事务回滚
        }
    }
    public static String truncateTitle(String title, int maxLength) {
        // 如果标题为空或长度不超过maxLength，直接返回
        if (title == null || title.length() <= maxLength) {
            return title;
        }
        // 截取前maxLength个字符
        String truncated = title.substring(0, maxLength);
        // 检查最后一个字符是否是空格，如果不是，则继续截断直到遇到空格
        if (!Character.isWhitespace(title.charAt(maxLength)) || !".".equals(title.charAt(maxLength))) {
            int lastSpaceIndex = truncated.lastIndexOf(' ');
            if (lastSpaceIndex != -1) {
                truncated = truncated.substring(0, lastSpaceIndex);
            }
            truncated = truncated + ".";
        }
        return truncated;
    }

} 